import torch
from tqdm import tqdm
import os


def fit_one_epoch(model_train, model, cross_loss, optimizer, epoch, epoch_step,
                  epoch_step_val, gen_train, gen_val, Epoch, cuda, local_rank, save_dir):
    loss = 0
    val_loss = 0
    if local_rank == 0:
        print('Start Train')
        pbar = tqdm(total=epoch_step, desc=f'Epoch {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3)

    model_train.train()
    for iteration, batch in enumerate(gen_train):
        if iteration >= epoch_step:
            break
        image, label = batch[0], batch[1]
        with torch.no_grad():
            if cuda:
                image = image.cuda()
                label = label.cuda()

        optimizer.zero_grad()
        # ----------------------#
        #   前向传播
        # ----------------------#
        outputs = model_train(image)

        # ----------------------#
        #   计算损失
        # ----------------------#
        loss_value = cross_loss(outputs, label)

        loss_value.backward()
        optimizer.step()
        loss += loss_value.item()

        if local_rank == 0:
            pbar.set_postfix(**{'loss': loss / (iteration + 1),
                                'lr': get_lr(optimizer)})
            pbar.update(1)

    if local_rank == 0:
        pbar.close()
        print('Finish Train')
        print('Start Validation')
        pbar = tqdm(total=epoch_step_val, desc=f'Epoch {epoch + 1}/{Epoch}', postfix=dict, mininterval=0.3)

    # 验证集验证
    model_train.eval()
    # 校验代码：不涉及参数训练
    for iteration, batch in enumerate(gen_val):
        if iteration >= epoch_step_val:
            break
        image, label = batch[0], batch[1]
        # 验证集校验 参数不变化
        with torch.no_grad():
            if cuda:
                image = image.cuda()
                label = label.cuda()
        # ----------------------#
        #   清零梯度
        # ----------------------#
        optimizer.zero_grad()
        # ----------------------#
        #   前向传播
        # ----------------------#
        outputs = model_train(image)
        # ----------------------#
        #   计算损失
        # ----------------------#
        loss_value = cross_loss(outputs, label)
        val_loss += loss_value.item()

        if local_rank == 0:
            pbar.set_postfix(**{'val_loss': val_loss / (iteration + 1)})
            pbar.update(1)

    if local_rank == 0:
        pbar.close()
        print('Epoch:' + str(epoch + 1) + '/' + str(Epoch))
        print('Total Loss: %.3f || Val Loss: %.3f ' % (loss / epoch_step, val_loss / epoch_step_val))

        torch.save(model.state_dict(), os.path.join(save_dir, "last_epoch_weights.pth"))



# ---------------------------------------------------#
#   获得学习率
# ---------------------------------------------------#
def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']
